# -*- coding: utf-8 -*-
"""
===============================
@Author     : Zuo WenTao
@Time       : 2024/7/10 8:41
@Description: 基于pymilvus2.4.x 版本编写
@Change     : 
@File       : MilvusClient.py
===============================
"""
from typing import List, Dict, Union
from pymilvus.client import utils
import pandas as pd
from pymilvus import (
    connections,
    CollectionSchema,
    Collection,
    FieldSchema,
    DataType,
    list_collections,
)
from loguru import logger


class MilvusConnectionError(Exception):
    """自定义 Milvus 连接错误异常。"""
    pass


class MilvusClient:

    def __init__(self, milvus_connect: Dict):
        """
        Milvus 向量数据库连接
        Args:
            table: 集合名称
            table_desc:  集合简介
            milvus_connect: 连接信息: {"host": xxx,"port": xxx,"user": xxx,"password": xxx}  其中 user以及password是可选参数
        """
        self.coll = None
        try:
            self._connect_to_milvus(milvus_connect)
            logger.success(f"[Milvus] --- 连接成功：{milvus_connect['host']}:{milvus_connect['port']}")
        except Exception as e:
            raise MilvusConnectionError(f"[Milvus] --- 连接失败：{e}")

    def _connect_to_milvus(self, milvus_connect: Dict):
        """连接到 Milvus 服务器。"""
        host = milvus_connect.get("host")
        port = milvus_connect.get("port")

        if not host or not port:
            raise ValueError("[Milvus] --- milvus_connect 缺少 'host' 或 'port'")

        connections.connect(**milvus_connect)
        logger.info(f"[Milvus] --- 连接到 Milvus：{host}:{port}")

    def create_collection(self, name: str, description: str = "", fields=None, emb_field_name: str = "embedding", is_index_emb=1):
        """
        Milvus 创建表以及向量索引
        Args:
            name: 要创建向量集合名称
            description: 向量集合介绍
            fields: 创建向量集合的字段列表,默认为only_id,embedding两个字段
            emb_field_name: 向量字段名称,用于创建向量字段索引,默认为 embedding
            is_index_emb: 是否创建向量字段索引,默认创建

        Returns:

        """
        if fields is None:
            fields = []
        if name in list_collections():
            self.coll = Collection(name)
        else:
            fields = fields or [
                FieldSchema(name="only_id", dtype=DataType.VARCHAR, max_length=100, is_primary=True, auto_id=False),
                FieldSchema(name="embedding", dtype=DataType.FLOAT_VECTOR, dim=1024),
            ]
            schema = CollectionSchema(fields, description)
            self.coll = Collection(name, schema, shards_num=4)
            if is_index_emb == 1:
                index_params = {
                    'metric_type': 'IP',
                    'index_type': 'HNSW',
                    'params': {'M': 64, 'efConstruction': 256}
                }
                self.coll.create_index(field_name=emb_field_name, index_params=index_params)
            logger.info(f"[Milvus] --- 创建集合 '{name}' 成功，描述：{description}")

    def create_index(self, index_items: Dict):
        """
        创建索引
        标量索引创建可以参考: https://milvus.io/docs/index-scalar-fields.md
        Args:
            index_items(Dict): {field_name: {index_param}}

        Returns:

        """
        for field_name, index_param in index_items.items():
            if isinstance(index_param, dict):
                self.coll.create_index(field_name=field_name, index_params=index_param)
            else:
                logger.info(f"[Milvus] --- 创建字段: {field_name} 的索引失败: index_param必须为字典")

    def bath_insert_data(self, data: Union[List, pd.DataFrame, Dict, utils.SparseMatrixInputType], timeout=120) -> bool:
        """
        向 Milvus 向量数据库中插入数据
        Args:
            data(list/tuple/pandas.DataFrame/sparse types): 需要传递一个列表, 常用传递格式: [[],[]],或者[{},{}]
                    注意: 列表中嵌套的列表顺序需要同 创建索引时候的字段顺序相同
        Returns: True or False

        """
        for x in range(3):
            try:
                self.coll.insert(data=data, timeout=timeout)
                return True
            except Exception as e:
                logger.exception(f"[Milvus] --- 向量写入失败 -- 失败原因: {e}")
        return False

    # def delete_data(self, id_name: str, delete_ids: List, type=1, timeout=120) -> bool:
    #     """
    #
    #     Args:
    #         delete_ids: 需要删除的向量的ID列表
    #         type: 1 按照筛选器进行删除 , 其余: 按照 id 删除实体(未找到可以进行删除的pymilvus版本)
    #
    #     Returns:
    #
    #     """
    #     try:
    #         if type == 1:
    #             self.coll.delete(expr=f"{id_name} in {str(delete_ids)}", timeout=timeout)
    #             return True
    #         else:
    #             # self.coll.delete(ids=delete_ids)
    #             return False
    #     except Exception as e:
    #         logger.error(e)
    #         return False

    def delete_data(self, expr: str, type=1, timeout=120) -> bool:
        """

        Args:
            delete_ids: 需要删除的向量的ID列表
            type: 1 按照筛选器进行删除 ,其余: 按照 id 删除实体(未找到可以进行删除的pymilvus版本)

        Returns:

        """
        try:
            if type == 1:
                self.coll.delete(expr=expr, timeout=timeout)
                return True
            else:
                # self.coll.delete(ids=delete_ids)
                return False
        except Exception as e:
            logger.error(e)
            return False

    def update_data(self, data: Union[List, pd.DataFrame, Dict, utils.SparseMatrixInputType], timeout=120):
        """
        更新 Milvus 向量数据库中的数据
        更新插入操作不会更新主键。
        禁止使用upsert方法代替大规模数据插入,会导致 Milvus 数据节点的内存消耗过高。
        Args:
            data(list/tuple/pandas.DataFrame/sparse types): 需要传递一个列表, 常用传递格式: [[],[]],或者[{},{}]
                    注意: 列表中嵌套的列表顺序需要同 创建索引时候的字段顺序相同
        Returns: True or False
        """
        for x in range(3):
            try:
                self.coll.upsert(data=data, timeout=timeout)
                return True
            except Exception as e:
                logger.exception(f"[Milvus] --- 向量更新失败 -- 失败原因: {e}")
        return False

    def search_by_embedding(self,
                            emb_data: List[List],
                            search_field: str,
                            out_field: List[str],
                            limit: int = 10,
                            expr=None
                            ) -> List:
        """
        向量检索
        :param emb_data:用于搜索的数据  list[list[Float]]
        :param search_field :要搜索的向量字段的名称
        :param out_field: 要输出的字段名称列表。指定后，可以使用 获取指定字段的值。
        :param limit: 返回的数据量
        :param expr: 过滤表达式
        :return:
        """
        search_param = {
            "data": emb_data,
            "anns_field": search_field,
            "param": {"metric_type": "IP", "params": {"nprobe": 10}},
            "offset": 0,
            "limit": limit,
            "expr": expr,
            "output_fields": out_field,

        }
        results = self.coll.search(**search_param)
        res_list = []
        for datas in results:
            for data in datas:
                res_list.append({
                    "Unique": data.id,
                    "distance": data.distance,
                    "fields": data.fields
                })
        return res_list

    def search_by_filter(
            self,
            expr: str,
            output_field: list = None,
            limit: int = 10
    ) -> List:
        """
        使用 Milvus 基于字段条件检索文档

        参数:
        expr -- 过滤条件表达式
        output_fields -- 要返回的字段列表，默认为 None 表示返回所有字段
        limit -- 返回的结果数量限制，默认值为 10

        返回:
        检索结果
        """
        results = self.coll.query(
            expr=expr,
            output_fields=output_field,
            limit=limit
        )
        res_list = []
        for datas in results:
            res_list.append(datas)

        return res_list

    def load(self):
        """用于向量上传完成之后,将表进行加载操作,只有加载成功, 创建的Milvus集合才可以进行检索"""
        self.coll.load()

    def flush_coll(self, timeout=60):
        """
        进行Milvus向量集合的更改操作,需要对集合中的数据从内存中进行Flush操作, 使数据进行落盘操作
        Milvus会自动进行落盘, 其会每一秒对现存集合的数据进行落盘操作
        使用delete操作,也可以调用Flush操作, 保证新增的数据可见,被删除的数据不会再被搜索到
        Args:
            timeout:
                    设置超时时间, 默认为60秒
        Returns:

        """
        self.coll.flush(timeout=timeout)


if __name__ == '__main__':
    table = "arxiv_test"
    table_desc = "封装工具类测试"
    milvus_connect = {"host": "192.168.200.59", "port": 32074}
    milvus_client = MilvusClient(milvus_connect)
    milvus_client.create_collection(table, table_desc)
    # expr = 'organ_code == "ENTRALIB001"'
    # expr = 'extend == "否"'
    # organ_codes = ["QMDXYXTS001","YKQXHWB","AHASARAK001","HONAENNI001","ENGUETTA001","EAALLEOL001","ENTRALIB001","INGONGKU001","YARIFIDA001","NIVERSIT007"]
    # for organ_code in organ_codes:
    #     expr = f'organ_code == "{organ_code}"'
    #     output_field = ['content']
    #     limit = 10000
    #     results = milvus_client.search_by_filter(expr, output_field=output_field, limit=limit)
    #     with open(rf"D:\yingke\data\chatlibrary_corpus_Thailand\{organ_code}.txt", "a", encoding="utf-8") as f:
    #         for result in results:
    #             f.write(result["content"] + "\n")



    expr = 'pdf_code == "arXiv:2401.00089"'
    output_field = ['embedding',"only_id","page","pdf_code"]
    limit = 10000
    results = milvus_client.search_by_filter(expr, output_field=output_field, limit=limit)

    # with open(f"{table}.txt", "a", encoding="utf-8") as f:
    #     for result in results:
    #         f.write(result["content"] + "\n")
"""
"QMDXYXTS001","YKQXHWB","AHASARAK001","HONAENNI001","ENGUETTA001","EAALLEOL001","ENTRALIB001","INGONGKU001","YARIFIDA001","NIVERSIT007"
"""